# 导入包
from matplotlib import pyplot as plt
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
import statsmodels.api as sm
import numpy as np
import seaborn as sns

# 数据库配置
db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': 'sjk1234',
    'database': 'tu',
    'port': 3306,
    'charset': 'sys'
}

# 创建数据库连接
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)
conn = pymysql.connect(**db_config)

# 从数据库读取数据
chunk_size = 10000
df = pd.read_sql_query(
    """
    SELECT d.*, m.buy_lg_vol, m.sell_lg_vol, m.buy_elg_vol, m.sell_elg_vol, m.net_mf_vol
    FROM date_1 d
    JOIN moneyflows m 
    ON d.ts_code = m.ts_code AND d.trade_date = m.trade_date
    LEFT JOIN index_daily i ON d.trade_date = i.trade_date AND i.ts_code = '399001.SZ'
    WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code = '000001.SZ'
    """,
    conn,
    chunksize=chunk_size
)
df1 = pd.concat(df, ignore_index=True)

# 数据预处理
df1['zd_closes'] = df1['zd_closes'].shift(1)
df1['zs_closes'] = round((df1['closes'] - df1['closes'].shift(1)) / df1['closes'].shift(1), 2)

# 处理缺失值
df1 = df1.dropna(subset=['zd_closes', 'zs_closes']).reset_index(drop=True)

# 选择数值型列
numeric_cols = df1.select_dtypes(include=['number']).columns.tolist()
print(df1.head())

# 特征选择
X = df1[['buy_lg_vol', 'sell_lg_vol', 'sell_elg_vol', 'net_mf_vol', 'zs_closes']]

# PCA分析
eigenvalues, eigenvectors = np.linalg.eig(np.cov(X, rowvar=False))
print('累计贡献率为:', round(eigenvalues[:5].sum() / eigenvalues.sum(), 4) * 100)

n_components = 5
top_eigenvectors = eigenvectors[:, :n_components]
principal_components = np.dot(X, top_eigenvectors)

# 创建PCA特征DataFrame
data_pca = pd.concat([df1, pd.DataFrame(principal_components,
                                      columns=[f'PC{i+1}' for i in range(n_components)])], axis=1)

# 回归分析
X_pca = data_pca[[f'PC{i+1}' for i in range(n_components)]].copy()
X_pca = sm.add_constant(X_pca)
y = df1['zd_closes'].copy()

model_pca = sm.OLS(y, X_pca)
result_pca = model_pca.fit()

print("\n回归模型结果:")
print(result_pca.summary())

# 选择前3个主成分
X_pca_selected = data_pca[['PC1', 'PC2', 'PC3']].copy()
X_pca_selected.columns = ['PC1', 'PC2', 'PC3']
X_pca_selected = sm.add_constant(X_pca_selected)

model_pca_selected = sm.OLS(y, X_pca_selected)
result_pca_selected = model_pca_selected.fit()

print("\n回归模型结果:")
print(result_pca_selected.summary())

# 绘制散点图
fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5))
for i, col in enumerate(['PC1', 'PC2', 'PC3']):
    axes[i].scatter(X_pca_selected[col], y, s=50, alpha=0.7)
    axes[i].set_xlabel(col)
    axes[i].set_ylabel("zd_closes")
    axes[i].set_title(f'{col} vs zd_closes')
plt.tight_layout()
plt.show()

# 输出主成分表达式
for k in range(5):
    stringy = f'PC{k+1} = '
    for j in range(len(eigenvectors[k])):
        coef = round(eigenvectors[k][j], 2)
        var_name = X.columns[j]
        if j > 0 and coef >= 0:
            stringy += ' + '
        elif coef < 0:
            stringy += ' - '
            coef = abs(coef)
        stringy += f'{coef}*{var_name}'
    print(stringy)